package vip.zhenzicheng.algorithm.leetcode.sort;

import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * <a href="https://leetcode.cn/problems/top-k-frequent-elements/">前K个高频元素 [中等]</a>
 * 与第K大的元素类似，参考{@link KthLargestElementInAnArray_215}
 *
 * @author zhenzicheng
 * @date 2022-06-20 14:05
 */
public class TopKFrequentElements_347 {

  /*使用JDK提供的优先级队列(堆排序)*/
  public int[] topKFrequent(int[] nums, int k) {
    // key保存元素，value保存出现的次数
    Map<Integer, Integer> map = new HashMap<>();
    for (int num : nums) {
      int count = map.getOrDefault(num, 0) + 1;
      map.put(num, count); // 替换原来的次数
    }

    PriorityQueue<int[]> minHeap = new PriorityQueue<>(k, (a, b) -> a[1] - b[1]); // 使用次数作为最小堆的比较规则
    for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
      int count = entry.getValue();
      if (minHeap.size() == k) {
        if (count > minHeap.peek()[1]) {
          minHeap.poll();
          minHeap.offer(new int[]{entry.getKey(), count});
        }
      } else {
        minHeap.offer(new int[]{entry.getKey(), count});
      }
    }

    // 输出
    int[] result = new int[k];
    for (int i = 0; i < k; i++) {
      result[i] = minHeap.poll()[0];
    }
    return result;
  }
}
